import { TreeNode } from '.'

// 统计树中所有节点的总数
const countTotalNodes = (nodes: TreeNode[]): number => {
  return nodes.reduce((count, node) => {
    const childrenCount = node.children ? countTotalNodes(node.children) : 0
    return count + 1 + childrenCount
  }, 0)
}

// 统计当前选中的节点数
const countCheckedNodes = (nodes: TreeNode[], checkedMap: Record<string, boolean>): number => {
  return nodes.reduce((count, node) => {
    const isChecked = !!checkedMap[node.id]
    const childrenCount = node.children ? countCheckedNodes(node.children, checkedMap) : 0
    return count + (isChecked ? 1 : 0) + childrenCount
  }, 0)
}

// 判断选中状态
export const getSelectionStatus = (
  nodes: TreeNode[],
  checkedMap: Record<string, boolean>
): 'none' | 'partial' | 'all' => {
  const totalNodeCount = countTotalNodes(nodes)
  const checkedNodeCount = countCheckedNodes(nodes, checkedMap)

  if (checkedNodeCount === 0) {
    return 'none' // 全未选中
  }
  if (checkedNodeCount === totalNodeCount) {
    return 'all' // 全部选中
  }
  return 'partial' // 部分选中
}

// 扁平化处理选中的树节点，返回 TreeNode[]
export const flatCheckedTreeNode = (checkedMap: Record<string, boolean>, data: TreeNode[]): TreeNode[] => {
  const result: TreeNode[] = []
  const addedNodes = new Set<string>() // 用于去重

  // 查找节点的所有子节点
  const collectChildren = (nodes: TreeNode[]) => {
    nodes.forEach((node) => {
      if (!addedNodes.has(node.id)) {
        result.push(node)
        addedNodes.add(node.id)
      }
      if (node.children) {
        collectChildren(node.children)
      }
    })
  }

  // 遍历 checkedMap，收集选中节点及其子节点
  Object.keys(checkedMap).forEach((id) => {
    if (checkedMap[id]) {
      const node = findNodeById(data, id)
      if (node) {
        collectChildren([node])
      }
    }
  })

  return result
}

// 查找指定 ID 的节点
const findNodeById = (nodes: TreeNode[], id: string): TreeNode | undefined => {
  for (const node of nodes) {
    if (node.id === id) {
      return node
    }
    if (node.children) {
      const found = findNodeById(node.children, id)
      if (found) return found
    }
  }
  return undefined
}
